import os

from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

from src.process_text import delete_bracket

try:
    from dotenv import load_dotenv
except ImportError:
    raise RuntimeError('Python environment for SPARK AI is not completely set up: required package "python-dotenv" is missing.') from None

load_dotenv()

def call_sparkai_model(prompt):

    spark = ChatSparkLLM(
        spark_api_url=os.environ["SPARKAI_URL"],
        spark_app_id=os.environ["SPARKAI_APP_ID"],
        spark_api_key=os.environ["SPARKAI_API_KEY"],
        spark_api_secret=os.environ["SPARKAI_API_SECRET"],
        spark_llm_domain=os.environ["SPARKAI_DOMAIN"],
        streaming=False,
    )
    response = get_response(prompt, spark)
    final_response = delete_bracket(response)
    return final_response

def get_response(prompt, spark):
    messages = [ChatMessage(
        role="user",
        content=prompt
    )]

    message = spark.generate([messages])  # return a list
    return message.generations[0][0].text